#include <map>
#include <string>
#include <vector>

#define PY_SSIZE_T_CLEAN
#define PY_ARRAY_UNIQUE_SYMBOL COIN2_ARRAY_API

#include <Python.h>
#include <numpy/arrayobject.h>
#include <pybind11/pybind11.h>

using namespace pybind11::literals;

class DataFrame {
 public:
  void AddIntCol(const std::string& name, const std::vector<int64_t> vec) {
    int_map_[name] = std::move(vec);
  }
  void AddDoubleCol(const std::string& name, const std::vector<double> vec) {
    double_map_[name] = std::move(vec);
  }
  std::vector<int64_t>& CreateIntCol(const std::string& name) {
    return int_map_[name] = std::vector<int64_t>();
  }
  std::vector<double>& CreateDoubleCol(const std::string& name) {
    return double_map_[name] = std::vector<double>();
  }

  size_t NumRows() {
    if (!int_map_.empty()) {
      return int_map_.begin()->second.size();
    }
    if (!double_map_.empty()) {
      return double_map_.begin()->second.size();
    }
    return 0;
  }
  size_t NumCols() { return int_map_.size() + double_map_.size(); }
  const std::map<std::string, std::vector<int64_t>>& IntCols() { return int_map_; }
  const std::map<std::string, std::vector<double>>& DoubleCols() { return double_map_; }

  std::vector<int64_t>& IntCol(const std::string& name) { return int_map_[name]; }
  std::vector<double>& DoubleCol(const std::string& name) { return double_map_[name]; }

 private:
  std::map<std::string, std::vector<int64_t>> int_map_;    // timestamp
  std::map<std::string, std::vector<double>> double_map_;  // most timeseries
};

// returned numpy.array cannot outlive vec
template <typename T>
pybind11::object vec_to_numpy(const std::vector<T>& vec) {
  npy_intp dim = vec.size();
  PyObject* ndarr = PyArray_NewFromDescr(
      &PyArray_Type,
      PyArray_DescrFromType((typeid(T) == typeid(int64_t)) ? NPY_INT64 : NPY_FLOAT64),
      1,                                                             // nd
      &dim,                                                          // dims. copied.
      nullptr,                                                       // strides
      const_cast<void*>(reinterpret_cast<const void*>(vec.data())),  // data
      NPY_ARRAY_CARRAY,                                              // flags
      nullptr                                                        // obj
  );

  return pybind11::reinterpret_borrow<pybind11::object>(ndarr);
}

/// lifetime
/// 1. A python script stops using PyDataFrame
/// 2. PyDataFrame is dealloced
/// 3. pd_ is dealloced (before df_)
/// 4. df_ is dealloced
class PyDataFrame {
 public:
  explicit PyDataFrame(std::shared_ptr<DataFrame> df) : df_(df) {
    pybind11::dict dict;
    num_rows_ = df_->NumRows();

    // 1. add ndarrays
    for (const auto& it : df_->IntCols()) {
      const auto& col_name = it.first;
      const auto col_array = vec_to_numpy(it.second);
      dict[col_name.c_str()] = col_array;
    }
    for (const auto& it : df_->DoubleCols()) {
      const auto& col_name = it.first;
      const auto col_array = vec_to_numpy(it.second);
      dict[col_name.c_str()] = col_array;
    }

    // 2. construct pandas.DataFram
    pd_ = pybind11::module::import("pandas").attr("DataFrame")(dict, "copy"_a = false);
  }

  pybind11::object df() { return pd_; }

 private:
  npy_intp num_rows_;
  std::shared_ptr<DataFrame> df_;
  pybind11::object pd_;  // pandas dataframe
};

PyDataFrame run() {
  std::shared_ptr<DataFrame> df(new DataFrame());

  auto& vec = df->CreateIntCol("col999");
  vec.push_back(41);
  vec.push_back(39);
  vec.push_back(50);
  return PyDataFrame(df);
}

int init_numpy() { import_array(); }

PYBIND11_MODULE(myexp, m) {
  init_numpy();
  m.doc() = "pybind11 example plugin";  // optional module docstring
  m.def("run", &run, "Run an exp");
  pybind11::class_<PyDataFrame>(m, "PyDataFrame").def("df", &PyDataFrame::df);
}
